#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GFPGAN图片处理器
支持COS URL工作流 + 自动清理
"""

import os
import time
import logging
import cv2
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
from typing import Dict, Optional
from config import Config
from cos_uploader import cos_handler
from cleanup_service import cleanup_service

logger = logging.getLogger(__name__)

class GFPGANProcessor:
    def __init__(self):
        self.task_status = {}  # 存储任务状态
        self.gpu_available = self._check_gpu()
        
        logger.info(f"GFPGAN处理器初始化完成，GPU可用: {self.gpu_available}")
    
    def _check_gpu(self) -> bool:
        """检查GPU是否可用"""
        try:
            import torch
            available = torch.cuda.is_available()
            if available:
                logger.info(f"检测到GPU: {torch.cuda.get_device_name()}")
            return available
        except Exception as e:
            logger.warning(f"GPU检测失败: {e}，将使用CPU处理")
            return False
    
    def get_gpu_info(self) -> Dict:
        """获取GPU信息"""
        if self.gpu_available:
            try:
                import torch
                return {
                    'device_count': torch.cuda.device_count(),
                    'current_device': torch.cuda.current_device(),
                    'device_name': torch.cuda.get_device_name(),
                    'memory_total': torch.cuda.get_device_properties(0).total_memory,
                    'memory_allocated': torch.cuda.memory_allocated(),
                    'memory_cached': torch.cuda.memory_reserved()
                }
            except Exception as e:
                return {'error': f'无法获取GPU信息: {e}'}
        else:
            return {'gpu': False, 'mode': 'CPU'}
    
    def is_gpu_available(self) -> bool:
        """检查GPU是否可用"""
        return self.gpu_available
    
    def process_from_url(self, image_url: str, task_id: str, quality_level: str = 'high') -> Dict:
        """
        完整处理流程：下载 → 处理 → 上传 → 清理
        
        Args:
            image_url: 原图的COS URL
            task_id: 任务ID
            quality_level: 质量等级
            
        Returns:
            dict: {"success": bool, "url": str, "processing_time": float} 或 {"success": False, "error": str}
        """
        start_time = time.time()
        input_path = None
        output_path = None
        
        try:
            # 1. 从COS下载原图
            logger.info(f"[{task_id}] 开始处理，下载原图: {image_url}")
            input_path = os.path.join(Config.UPLOAD_FOLDER, f"{task_id}_original.jpg")
            
            if not cos_handler.download_from_url(image_url, input_path):
                return {'success': False, 'error': 'COS下载原图失败'}
            
            # 2. GFPGAN处理
            logger.info(f"[{task_id}] 开始图片增强处理")
            output_path = os.path.join(Config.RESULT_FOLDER, f"enhanced_{task_id}.jpg")
            
            success = self.enhance_image(input_path, output_path, quality_level=quality_level)
            
            if not success:
                return {'success': False, 'error': '图片处理失败'}
            
            # 3. 上传处理结果到COS
            logger.info(f"[{task_id}] 上传处理结果到COS")
            result_url = cos_handler.upload_file(output_path, task_id)
            
            if not result_url:
                return {'success': False, 'error': 'COS上传处理结果失败'}
            
            processing_time = time.time() - start_time
            logger.info(f"[{task_id}] 处理完成，耗时: {processing_time:.2f}s，结果URL: {result_url}")
            
            return {
                'success': True,
                'url': result_url,
                'processing_time': processing_time
            }
            
        except Exception as e:
            logger.error(f"[{task_id}] 处理异常: {e}")
            return {'success': False, 'error': str(e)}
            
        finally:
            # 4. ⭐ 无论成功失败，都清理本地文件
            cleanup_service.cleanup_task_files(task_id)
    
    def enhance_image(self, input_path: str, output_path: str, quality_level: str = 'high') -> bool:
        """
        图片增强处理
        根据文件大小自动选择处理方式
        
        Args:
            input_path: 输入图片路径
            output_path: 输出图片路径
            quality_level: 质量等级 (low/medium/high)
            
        Returns:
            bool: 处理是否成功
        """
        try:
            file_size = os.path.getsize(input_path)
            logger.info(f"图片大小: {file_size / 1024:.2f} KB")
            
            # 根据文件大小选择处理方式
            if file_size > 500000:  # 大于500KB使用OpenCV
                logger.info("使用OpenCV进行高质量处理")
                result = self._process_with_opencv(input_path, output_path, quality_level)
            else:
                logger.info("使用PIL进行快速处理")
                result = self._process_with_pil(input_path, output_path, quality_level)
            
            return result['success']
            
        except Exception as e:
            logger.error(f"图片增强异常: {e}")
            return False
    
    def _process_with_opencv(self, input_path: str, output_path: str, quality_level: str) -> Dict:
        """使用OpenCV进行图片增强处理"""
        try:
            # 读取输入图片
            img = cv2.imread(input_path, cv2.IMREAD_COLOR)
            if img is None:
                raise ValueError(f"无法读取图片: {input_path}")
            
            # 1. 去噪处理
            denoised = cv2.fastNlMeansDenoisingColored(img, None, 10, 10, 7, 21)
            
            # 2. 锐化处理
            kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
            sharpened = cv2.filter2D(denoised, -1, kernel)
            
            # 3. 对比度增强
            lab = cv2.cvtColor(sharpened, cv2.COLOR_BGR2LAB)
            l, a, b = cv2.split(lab)
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
            l = clahe.apply(l)
            enhanced = cv2.merge([l, a, b])
            enhanced = cv2.cvtColor(enhanced, cv2.COLOR_LAB2BGR)
            
            # 4. 色彩饱和度增强
            hsv = cv2.cvtColor(enhanced, cv2.COLOR_BGR2HSV)
            hsv[:,:,1] = hsv[:,:,1] * 1.2
            hsv[:,:,1] = np.clip(hsv[:,:,1], 0, 255)
            enhanced = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
            
            # 5. 边缘增强
            gray = cv2.cvtColor(enhanced, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, 50, 150)
            edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
            enhanced = cv2.addWeighted(enhanced, 0.8, edges_colored, 0.2, 0)
            
            # 保存增强后的图片
            cv2.imwrite(output_path, enhanced)
            
            output_size = os.path.getsize(output_path)
            logger.info(f"OpenCV处理完成，输出大小: {output_size / 1024:.2f} KB")
            
            return {'success': True, 'output_size': output_size}
            
        except Exception as e:
            logger.error(f"OpenCV处理失败: {e}")
            return {'success': False, 'error': str(e)}
    
    def _process_with_pil(self, input_path: str, output_path: str, quality_level: str) -> Dict:
        """使用PIL进行图片增强处理"""
        try:
            with Image.open(input_path) as img:
                # 1. 锐化处理
                sharpened = img.filter(ImageFilter.SHARPEN)
                
                # 2. 对比度增强
                enhancer = ImageEnhance.Contrast(sharpened)
                contrast_enhanced = enhancer.enhance(1.3)
                
                # 3. 亮度调整
                brightness_enhancer = ImageEnhance.Brightness(contrast_enhanced)
                brightness_enhanced = brightness_enhancer.enhance(1.1)
                
                # 4. 色彩饱和度增强
                color_enhancer = ImageEnhance.Color(brightness_enhanced)
                color_enhanced = color_enhancer.enhance(1.2)
                
                # 5. 保存增强后的图片
                color_enhanced.save(output_path, 'JPEG', quality=95, optimize=True)
            
            output_size = os.path.getsize(output_path)
            logger.info(f"PIL处理完成，输出大小: {output_size / 1024:.2f} KB")
            
            return {'success': True, 'output_size': output_size}
            
        except Exception as e:
            logger.error(f"PIL处理失败: {e}")
            return {'success': False, 'error': str(e)}
    
    # 任务状态管理方法（保留用于向后兼容）
    def init_task_status(self, task_id: str, status: str):
        """初始化任务状态"""
        self.task_status[task_id] = {
            'task_id': task_id,
            'status': status,
            'created_at': time.time(),
            'cos_url': None
        }
    
    def update_task_status(self, task_id: str, status: str, cos_url: Optional[str] = None):
        """更新任务状态"""
        if task_id in self.task_status:
            self.task_status[task_id]['status'] = status
            self.task_status[task_id]['updated_at'] = time.time()
            if cos_url:
                self.task_status[task_id]['cos_url'] = cos_url
    
    def get_task_status(self, task_id: str) -> Optional[Dict]:
        """获取任务状态"""
        return self.task_status.get(task_id)
